Write a custom CUDA kernel to optimize `CosFace Loss` (Large Margin Cosine Loss).

Formula: Loss = -log( exp(s * (cos_theta_yi - m)) / Sum(exp(s * cos_theta_j_modified)) )
Where:
- `cos_theta` is the cosine similarity matrix (Batch, Classes).
- `yi` is the ground truth class index for the sample.
- `m` is the additive cosine margin.
- `s` is the scaling factor.
- For the target class `j == yi`, the logit is `s * (cos_theta - m)`.
- For other classes `j != yi`, the logit is `s * cos_theta`.

Problem Analysis:
1. Memory Overhead: A standard implementation uses `torch.scatter` or `one_hot` multiplication to subtract `m` only from the target indices. This allocates auxiliary tensors equal to the size of the logits (N, C), wasting memory bandwidth.
2. Operator Chaining: The sequence `scatter/sub` -> `scale` -> `CrossEntropy (LogSoftmax -> NLL)` involves multiple kernel launches and redundant global memory reads/writes.

Optimization Strategy: Fused Logits-Modification and CrossEntropy

The strategy is to fuse the margin application, scaling, softmax normalization, and loss calculation into a single pass.

1. One-Block-per-Row: Each CUDA block processes one sample (one row of the cosine matrix) to calculate its loss.

2. On-the-Fly Logic:
   - Load cosine values from global memory.
   - Check if the current column index matches the target label `y_i`.
   - If match: val = s * (val - m).
   - If not match: val = s * val.
   This eliminates the need for one-hot masks or scatter operations.

3. Online Softmax (LogSumExp):
   - Use the 2-pass reduction algorithm (or 1-pass online algorithm) within the block to compute `Max` and `SumExp` of the *modified* logits.
   - This ensures numerical stability without materializing the full modified logits tensor.

4. Fused NLL Loss:
   - Calculate `log_prob = target_logit - (max_val + log(sum_exp))`.
   - Output `-log_prob`.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F

# 人脸识别场景：Batch Size 较小，但分类数 (Classes/Identities) 极大
BATCH_SIZE = 512
NUM_CLASSES = 10000 
SHAPE = (BATCH_SIZE, NUM_CLASSES)

# CosFace 超参数
SCALE_S = 64.0
MARGIN_M = 0.35

class CosFaceLoss(nn.Module):
    """
    Standard PyTorch implementation of CosFace Loss.
    """
    def __init__(self, s=64.0, m=0.35, reduction='mean'):
        super(CosFaceLoss, self).__init__()
        self.s = s
        self.m = m
        self.reduction = reduction

    def forward(self, cosine: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        # cosine: (N, C) - Normalized Features @ Normalized Weights
        # label: (N)
        
        # 1. 创建 One-hot 掩码 
        one_hot = torch.zeros_like(cosine)
        one_hot.scatter_(1, label.view(-1, 1), 1.0)
        
        # 2. Apply Margin: cos_theta - m (only for target class)
        logits = cosine - one_hot * self.m
        
        # 3. Scale: s * logits
        logits = logits * self.s
        
        # 4. CrossEntropy
        loss = F.cross_entropy(logits, label, reduction='none')
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, s=64.0, m=0.35, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = CosFaceLoss(s=s, m=m, reduction=reduction)
    
    def forward(self, cosine, label):
        return self.loss_fn(cosine, label)

def get_inputs():
    # 模拟归一化后的 Cosine Similarity (-1 ~ 1)
    cosine = torch.randn(SHAPE, dtype=torch.float32)
    # 归一化到 [-1, 1] 模拟真实余弦值
    cosine = torch.clamp(cosine, -1.0, 1.0)
    
    label = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [cosine.contiguous(), label.contiguous()]

def get_init_inputs():
    return [SCALE_S, MARGIN_M, 'none']